import acl
import cv2
import numpy as np

min_chn = np.array([123.675, 116.28, 103.53], dtype=np.float32)
var_reci_chn = np.array([0.0171247538316637, 0.0175070028011204, 0.0174291938997821], dtype=np.float32)
channel = 3
modelHeight = 256
modelWeight = 256
ACL_MEMCPY_DEVICE_TO_DEVICE = 3
ACL_MEM_MALLOC_HUGE_FIRST = 0
classNum = 30


def ProcessInput(input_path):
    image = cv2.imread(input_path)
    resize_image = cv2.resize(image, (modelWeight, modelHeight))

    # switch channel BGR to RGB
    image_rgb = cv2.cvtColor(resize_image, cv2.COLOR_BGR2RGB)
    image_rgb = image_rgb.astype(np.float32)

    # data standardization
    image_rgb = (image_rgb - min_chn) * var_reci_chn

    # HWC to CHW
    image_rgb = image_rgb.transpose([2, 0, 1]).copy()
    return image_rgb.tobytes()

def main():
    # ***************************** 初始准备工作 *****************************
    acl.init()              # 初始化
    acl.rt.set_device(0)    # 设置使用第1块NPU
    modelId, ret = acl.mdl.load_from_file("./resnet50.om")  # 加载模型



    # ***************************** 准备模型输入 *****************************
    # 对输入图片进行处理
    cpu_input_buffer = ProcessInput("2.jpg")

    # 将cpu上数据拷贝到npu上
    input_data_size = len(cpu_input_buffer)
    npu_input_buffer, ret = acl.rt.malloc(input_data_size, ACL_MEM_MALLOC_HUGE_FIRST)
    cpu_input_buffer = acl.util.bytes_to_ptr(cpu_input_buffer)
    ret = acl.rt.memcpy(npu_input_buffer, input_data_size, cpu_input_buffer, input_data_size, ACL_MEMCPY_DEVICE_TO_DEVICE)
    # 把NPU中分配的显存绑定到aclmdlDataset结构体上
    inputData = acl.create_data_buffer(npu_input_buffer, input_data_size)
    inputDataset = acl.mdl.create_dataset()
    inputDataset, ret = acl.mdl.add_dataset_buffer(inputDataset, inputData)



    # ***************************** 准备模型输出空间 *****************************
    output_buffer_size = classNum * 4
    npu_output_buffer, ret = acl.rt.malloc(output_buffer_size, ACL_MEM_MALLOC_HUGE_FIRST)
    outputData = acl.create_data_buffer(npu_output_buffer, output_buffer_size)
    outputDataset = acl.mdl.create_dataset()
    outputDataset, ret = acl.mdl.add_dataset_buffer(outputDataset, outputData)



    # ***************************** 模型执行推理 *****************************
    ret = acl.mdl.execute(modelId, inputDataset, outputDataset)



    # ***************************** 模型输出后处理 *****************************
    out_class_data, ret = acl.rt.malloc(output_buffer_size, ACL_MEM_MALLOC_HUGE_FIRST)
    ret = acl.rt.memcpy(out_class_data, output_buffer_size, npu_output_buffer, output_buffer_size, ACL_MEMCPY_DEVICE_TO_DEVICE)
    
    # printMax(out_class_data)函数部分
    bytes_data = acl.util.ptr_to_bytes(out_class_data, output_buffer_size)
    data = np.frombuffer(bytes_data, dtype=np.float32).reshape([1, 30])
    data = data.flatten()
    vals = np.exp(data)/np.sum(np.exp(data))
    top_index = vals.argsort()[-1]
    print(f"class:{top_index} conf:{vals[top_index]:06f}")



    # ***************************** 释放昇腾资源 *****************************
    acl.rt.free(npu_input_buffer)
    acl.mdl.destroy_dataset(inputDataset)
    acl.rt.free(npu_output_buffer)
    acl.mdl.destroy_dataset(outputDataset)
    ret = acl.mdl.unload(modelId)
    acl.rt.reset_device(0)
    acl.finalize()



if __name__ == "__main__":
    main()